# # ChatDoc:又一个智能文档助手
#
# - 读取pdf、excel、doc三种常见的文档格式
# - 根据文档内容，智能抽取内容并输出相应格式

# 整合优化，动态加载三种文件格式,增加了文本切割

#导入必须的包
from langchain_community.document_loaders import UnstructuredExcelLoader,Docx2txtLoader,PyPDFLoader
from langchain.text_splitter import  CharacterTextSplitter


#定义chatdoc
class ChatDoc():

    def __init__(self):
        self.doc = None
        self.splitText = [] #分割后的文本

    def getFile(self):
        doc = self.doc
        loaders = {
            "docx":Docx2txtLoader,
            "pdf":PyPDFLoader,
            "xlsx":UnstructuredExcelLoader,
        }
        file_extension = doc.split(".")[-1]
        loader_class = loaders.get(file_extension)
        if loader_class:
            try:
                loader = loader_class(doc)
                text = loader.load()
                return text
            except Exception as e:
                print(f"Error loading {file_extension} files:{e}")
        else:
             print(f"Unsupported file extension: {file_extension}")
             return  None

    #处理文档的函数
    def splitSentences(self):
        full_text = self.getFile() #获取文档内容
        if full_text != None:
            #对文档进行分割
            text_split = CharacterTextSplitter(
                chunk_size=150,
                chunk_overlap=20,
            )
            texts = text_split.split_documents(full_text)
            self.splitText = texts

chat_doc = ChatDoc()
chat_doc.doc = "example/fake.xlsx"
# 没有返回结果单纯分词
chat_doc.splitSentences()
print(chat_doc.splitText)